% 清除工作区和命令窗口
clear all;
close all;
clc;

% 参数设置
numSamples = 10000; % 样本数量
numAntennas = 4; % 天线数量
numSubcarriers = 64; % 子载波数量
trainRatio = 0.8; % 训练集比例
SNRdBs = 0:5:30; % 不同的信噪比 (dB)
numSNRs = length(SNRdBs);

% 初始化性能指标数组
mseVals = zeros(numSNRs, 1);
rmseVals = zeros(numSNRs, 1);
psnrVals = zeros(numSNRs, 1);
ssimVals = zeros(numSNRs, 1);

for i = 1:numSNRs
    SNRdB = SNRdBs(i);
    
    % 生成通信数据
    H = generateChannelData(numSamples, numAntennas, numSubcarriers);
    Y = generateReceivedSignals(H, SNRdB);
    
    % 划分训练集和测试集
    trainSize = floor(trainRatio * numSamples);
    trainH = H(1:trainSize, :, :);
    trainY = Y(1:trainSize, :, :);
    testH = H(trainSize+1:end, :, :);
    testY = Y(trainSize+1:end, :, :);
    
    % 数据预处理
    trainY = reshape(trainY, [trainSize, numAntennas * numSubcarriers]);
    testY = reshape(testY, [numSamples - trainSize, numAntennas * numSubcarriers]);
    trainH = reshape(trainH, [trainSize, numAntennas * numSubcarriers]);
    testH = reshape(testH, [numSamples - trainSize, numAntennas * numSubcarriers]);
    
    % 归一化数据
    trainY = normalize(trainY);
    testY = normalize(testY);
    trainH = normalize(trainH);
    testH = normalize(testH);
    
    % 构建深度学习模型
    layers = [
        featureInputLayer(numAntennas * numSubcarriers, 'Name', 'input')
        fullyConnectedLayer(256, 'Name', 'fc1')
        reluLayer('Name', 'relu1')
        fullyConnectedLayer(128, 'Name', 'fc2')
        reluLayer('Name', 'relu2')
        fullyConnectedLayer(numAntennas * numSubcarriers, 'Name', 'output')
        regressionLayer('Name', 'regression')
    ];
    
    options = trainingOptions('adam', ...
        'MaxEpochs', 50, ...
        'MiniBatchSize', 64, ...
        'Shuffle', 'every-epoch', ...
        'Verbose', false, ...
        'Plots', 'training-progress');
    
    net = trainNetwork(trainY, trainH', layers, options);
    
    % 进行信道估计
    estimatedH = predict(net, testY');
    estimatedH = reshape(estimatedH', [numSamples - trainSize, numAntennas, numSubcarriers]);
    testH = reshape(testH, [numSamples - trainSize, numAntennas, numSubcarriers]);
    
    % 性能评估
    % 均方误差 (MSE)
    mse = mean(mean(mean((abs(estimatedH - testH)).^2, 3), 2), 1);
    mseVals(i) = mse;
    
    % 均方根误差 (RMSE)
    rmse = sqrt(mse);
    rmseVals(i) = rmse;
    
    % 峰值信噪比 (PSNR)
    maxVal = max(max(max(abs(testH))));
    psnr = 20 * log10(maxVal / rmse);
    psnrVals(i) = psnr;
    
    % 结构相似性指数 (SSIM)
    ssimSum = 0;
    for j = 1:size(testH, 1)
        ssimSum = ssimSum + ssim2(squeeze(abs(testH(j, :, :))), squeeze(abs(estimatedH(j, :, :))));
    end
    ssimVals(i) = ssimSum / size(testH, 1);
end

% 绘制性能指标随信噪比变化的曲线
figure;
subplot(2,2,1);
plot(SNRdBs, mseVals, 'b-o');
xlabel('信噪比 (dB)');
ylabel('均方误差 (MSE)');
title('MSE 随信噪比变化');
grid on;

subplot(2,2,2);
plot(SNRdBs, rmseVals, 'g-o');
xlabel('信噪比 (dB)');
ylabel('均方根误差 (RMSE)');
title('RMSE 随信噪比变化');
grid on;

subplot(2,2,3);
plot(SNRdBs, psnrVals, 'r-o');
xlabel('信噪比 (dB)');
ylabel('峰值信噪比 (PSNR)');
title('PSNR 随信噪比变化');
grid on;

subplot(2,2,4);
plot(SNRdBs, ssimVals, 'm-o');
xlabel('信噪比 (dB)');
ylabel('结构相似性指数 (SSIM)');
title('SSIM 随信噪比变化');
grid on;

% 生成信道数据的函数
function H = generateChannelData(numSamples, numAntennas, numSubcarriers)
    H = zeros(numSamples, numAntennas, numSubcarriers);
    for i = 1:numSamples
        for j = 1:numAntennas
            for k = 1:numSubcarriers
                H(i, j, k) = (randn + 1j*randn)/sqrt(2);
            end
        end
    end
end

% 生成接收信号的函数
function Y = generateReceivedSignals(H, SNRdB)
    numSamples = size(H, 1);
    numAntennas = size(H, 2);
    numSubcarriers = size(H, 3);
    Y = zeros(numSamples, numAntennas, numSubcarriers);
    for i = 1:numSamples
        for j = 1:numAntennas
            for k = 1:numSubcarriers
                s = (randn + 1j*randn)/sqrt(2); % 发送信号
                noisePower = 10^(-SNRdB/10);
                noise = (randn + 1j*randn)*sqrt(noisePower/2);
                Y(i, j, k) = H(i, j, k) * s + noise;
            end
        end
    end
end

% 归一化数据的函数
function normalizedData = normalize(data)
    meanData = mean(data, 1);
    stdData = std(data, 0, 1);
    normalizedData = (data - repmat(meanData, size(data, 1), 1))./repmat(stdData, size(data, 1), 1);
end
    